from flask import Flask, render_template, request, jsonify
import pandas as pd
import subprocess
import os
import sys

app = Flask(__name__, static_folder='js', static_url_path='/js')

initial_cash = 1000000
current_index = 0
portfolio = {'cash': initial_cash, 'position': 0, 'history': []}

def load_data():
    global data
    data = pd.read_csv('chinext_index.csv')
    data['timestamp'] = pd.to_datetime(data['timestamp'])

load_data()

@app.route('/')
def index():
    global current_index, portfolio
    # 重置数据
    current_index = 0
    portfolio = {'cash': initial_cash, 'position': 0, 'history': []}
    return render_template('index.html')

@app.route('/next_kline', methods=['POST'])
def next_kline():
    global current_index
    if current_index >= len(data):
        return jsonify({'status': 'end'})

    kline = data.iloc[current_index].to_dict()
    current_index += 1
    return jsonify(kline)

@app.route('/trade', methods=['POST'])
def trade():
    action = request.form.get('action')
    amount = float(request.form.get('amount'))
    price = data.iloc[current_index - 1]['close']  # 使用当前K线的收盘价进行交易
    if action == 'buy':
        cost = amount * price
        if portfolio['cash'] >= cost:
            portfolio['cash'] -= cost
            portfolio['position'] += amount
            portfolio['history'].append(f'Bought {amount} at {price}')
    elif action == 'sell':
        if portfolio['position'] >= amount:
            portfolio['cash'] += amount * price
            portfolio['position'] -= amount
            portfolio['history'].append(f'Sold {amount} at {price}')
    return jsonify(portfolio)

@app.route('/generate_random_klines', methods=['POST'])
def generate_random_klines():
    global current_index, portfolio
    try:
        script_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'chinext_index.csv.py')
        result = subprocess.run([sys.executable, script_path], check=True, capture_output=True, text=True, encoding='utf-8')
        # 重新加载数据
        load_data()
        # 重置数据
        current_index = 0
        portfolio = {'cash': initial_cash, 'position': 0, 'history': []}
        return jsonify({'status': 'success', 'output': result.stdout})
    except subprocess.CalledProcessError as e:
        return jsonify({'status': 'error', 'message': e.stderr or str(e)})

if __name__ == '__main__':
    app.run(debug=True)
